package com.example.mcp.config;

import org.springframework.http.HttpRequest;
import org.springframework.http.HttpStatus;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.stereotype.Component;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;

/**
 * OpenAI API响应拦截器
 * 在响应返回给Spring AI之前修复兼容性问题
 * 
 * @author Generated
 * @version 1.0.0
 */
@Component
public class OpenAiResponseInterceptor implements ClientHttpRequestInterceptor {

    private final OpenAiResponseProcessor responseProcessor;

    public OpenAiResponseInterceptor(OpenAiResponseProcessor responseProcessor) {
        this.responseProcessor = responseProcessor;
    }
    
    public OpenAiResponseProcessor getResponseProcessor() {
        return responseProcessor;
    }

    @Override
    public ClientHttpResponse intercept(
            HttpRequest request,
            byte[] body,
            ClientHttpRequestExecution execution) throws IOException {
        
        // 执行原始请求
        ClientHttpResponse response = execution.execute(request, body);
        
        // 只处理聊天完成请求
        String uri = request.getURI().toString();
        if (uri.contains("/chat/completions")) {
            return new ProcessedClientHttpResponse(response, responseProcessor);
        }
        
        return response;
    }

    /**
     * 自定义响应包装器
     */
    private static class ProcessedClientHttpResponse implements ClientHttpResponse {
        
        private final ClientHttpResponse originalResponse;
        private final OpenAiResponseProcessor processor;
        private byte[] processedBody;

        public ProcessedClientHttpResponse(ClientHttpResponse originalResponse, 
                                         OpenAiResponseProcessor processor) {
            this.originalResponse = originalResponse;
            this.processor = processor;
        }

        @Override
        public InputStream getBody() throws IOException {
            if (processedBody == null) {
                try {
                    // 读取原始响应
                    String originalBody = new String(originalResponse.getBody().readAllBytes(), 
                                                   StandardCharsets.UTF_8);
                    
                    // 处理响应内容
                    String processedBodyString = processor.processResponse(originalBody);
                    processedBody = processedBodyString.getBytes(StandardCharsets.UTF_8);
                } catch (IOException e) {
                    System.err.println("读取响应体时出错: " + e.getMessage());
                    throw e;
                }
            }
            
            return new ByteArrayInputStream(processedBody);
        }

        @Override
        public org.springframework.http.HttpHeaders getHeaders() {
            return originalResponse.getHeaders();
        }

        @Override
        public org.springframework.http.HttpStatusCode getStatusCode() {
            try {
                return originalResponse.getStatusCode();
            } catch (IOException e) {
                System.err.println("获取状态码时出错: " + e.getMessage());
                // 返回默认状态码
                return HttpStatus.OK;
            }
        }

        @Override
        public String getStatusText() throws IOException {
            return originalResponse.getStatusText();
        }

        @Override
        public void close() {
            originalResponse.close();
        }
    }
}
